"""
Created on 19 Dec 2020
@fake_author:JemusD

"""
from matplotlib.font_manager import FontProperties
from matplotlib import pyplot as plt
from math import log
import operator
# pickle包可以将决策树保存下来，方便下次直接调用
import pickle


"""
instructions: 计算信源熵 
    Sigma(p(xi)*log( p(xi)^(-1) ))
    p(xi):表示xi号标签出现的频率/标签总数
Parameters: 
    dataSet
Returns:
    shannonEnt - 信息熵
Modify:
    2020-12-19
"""
def calcShannonEnt(dataSet):
    numEntries = len(dataSet)
    labelCounts = {}
    for featVec in dataSet:
        currentLabels = featVec[-1]
        if currentLabels not in labelCounts.keys():
            # 字典里没有这个key-value对，就会自动创建一个
            labelCounts[currentLabels] = 0
        labelCounts[currentLabels] += 1
    # 计算熵
    shannonEnt = 0.0
    for key in labelCounts:
        prob = float(labelCounts[key]) / numEntries
        shannonEnt -= prob*log(prob, 2)
    
    return shannonEnt


"""
instructions: 创建一个简单的数据集
    用字典存储决策树结构：
    {'有自己的房子':{0:{'有工作':{0:'no', 1:'yes'}}, 1:'yes'}}
    年龄：0代表青年，1代表中年，2代表老年
    有工作：0代表否，1代表是
    有自己的房子：0代表否，1代表是
    信贷情况：0代表一般，1代表好，2代表非常好
    类别（是否给贷款）：no代表否，yes代表是
"""
def createDataSet():
    # 数据集
    dataSet = [[0, 0, 0, 0, 'no'],
               [0, 0, 0, 1, 'no'],
               [0, 1, 0, 1, 'yes'],
               [0, 1, 1, 0, 'yes'],
               [0, 0, 0, 0, 'no'],
               [1, 0, 0, 0, 'no'],
               [1, 0, 0, 1, 'no'],
               [1, 1, 1, 1, 'yes'],
               [1, 0, 1, 2, 'yes'],
               [1, 0, 1, 2, 'yes'],
               [2, 0, 1, 2, 'yes'],
               [2, 0, 1, 1, 'yes'],
               [2, 1, 0, 1, 'yes'],
               [2, 1, 0, 2, 'yes'],
               [2, 0, 0, 0, 'no']]
    # 分类属性
    labels = ['年龄', '有工作', '有自己的房子', '信贷情况']
    # 返回数据集和分类属性
    """
    dataSet = [[1, 1, 'A'],
                [1, 0, 'B'],
                [0, 1, 'A'],
                [0, 1, 'A']]
    labels = ['no surfacing', 'flippers']
    """
    return dataSet, labels


"""
instructions: 通过指定分类的列，和该列是否等于莫个value来划分数据集 
    及，返回的retDataset里面只包含featVec[axis]=value的那些数据集合
Parameters: 
    dataSet - 要分类的数据集
    axis - 通过该列来划分数据集
    value - 划分的值是否等于value来分开数据集
Returns:
    retDataSet - 返回满足要求的数据集
Modify:
    2020-12-19
"""
def splitDataSet(dataSet, axis, value):
    retDataSet = []
    for featVec in dataSet:
        if featVec[axis] == value:
            # rF 用来将该行中除第axis列以外其他列读取出来
            reducedFeatVec = featVec[:axis]
            reducedFeatVec.extend(featVec[axis+1:])
            # 将读取出来的行放入retDataSet
            retDataSet.append(reducedFeatVec)
    return retDataSet


"""
instructions: 选取收益最大的特征划分方式 
    1.获取特征维数
    2.计算总体数据集合的信源熵
    3.计算每一种划分方式的信源熵，然后减去总体的熵得到information gain。获得有最大information gain的
        及为最好的划分方式
    * 划分后的信源熵计算方式: len(part1)/total * H(part1) + len(part2)/total * H(part2)
        H(x) 表示x的信源熵
    
Parameters: 
    dataSet - 数据集
Returns:
    bestFeature - 按该列划分有最大信息收益, 一个i值
Modify:
    2020-12-19
"""
def chooseBestFeatureToSplit(dataSet):
    numFeatures = len(dataSet[0]) - 1
    baseEntries = calcShannonEnt(dataSet)
    bestInfoGain = 0.0
    bestFeature = -1
    for i in range(numFeatures):
        # 将第i列的特征全部取出放入featList
        featList = [example[i] for example in dataSet]
        uniqueVals = set(featList)
        newEntries = 0.0
        for value in uniqueVals:
            # 获取value分别取该特征的不同的类型的子集
            subDataSet = splitDataSet(dataSet, i, value)
            # 此处prob是各子数据集的权值
            prob = len(subDataSet)/float(len(dataSet))
            # 这步是类似:(3/5)*ent1 + (2/5)*ent2 的意思。每个子集的熵乘以权值。
            newEntries += prob*calcShannonEnt(subDataSet)
        infoGain = baseEntries - newEntries
        #print("the col of %d information gain is: %.3f" % (i, infoGain))
        if infoGain > bestInfoGain:
            bestInfoGain = infoGain
            bestFeature = i

    return bestFeature


"""
instructions: 将该类标签list中出现次数最多的返回
    # 当将所有标签labels分类完，如果同一个分类下还有不同种类，就以少数服从多数原则，指定多数为这个分类下的类型
Parameters: 
    classList - 标签list ['yes','no','yes','yes']
Returns:
    sortedClassCount - 
Modify:
    2020-12-20
"""
def majorityCnt(classList):
    classCount = {}
    for vote in classList:
        if vote not in classCount.keys(): 
            classCount[vote] = 0
        classCount[vote] += 1
    sortedClassCount = sorted(classCount.items(), key=operator.itemgetter(1), reverse=True)
    return sortedClassCount[0][0]


"""
instructions: 创建决策树
    bestFeat - 某个属性类，特征列的索引
    1.获取当前数据集的right answer列
    2.(终止条件)看该数据集是否元素都是同一类，或是该数据集已经没有属性列了
    3.获取划分最佳属性列，并记录其标签名字
    4.通过刚才的标签逐步创建决策树(涉及递归)
Parameters: 
     dataSet - 生成决策树的数据集
     labels - 属性列对应的名称标签 ['年龄', '有工作', '有自己的房子', '信贷情况']
     featLabels - 已经用来分过类的标签,这个是有先后顺序的。越是前面的标签，其信息增益在同样多属性分类下越大
Returns:
    myTree - 一个字典树(注意这里使用了递归的方式生成字典树) like this: {'you_jiji':{ 0:{}, 1:{} }}
Modify:
    2020-12-20
"""
def createTree(dataSet, labels, featLabels):
    classList = [example[-1] for example in dataSet]
    # 终止条件有两个1.已经遍历了所有特征属性。及len(dataSet[0]==1)
    #             2.每个分支下数据哦都有相同的分类
    # x.count(x[0]) 表示x中有多少个x[0]
    if classList.count(classList[0]) == len(classList):
        return classList[0]
    if len(dataSet[0]) == 1:
        # 当将所有标签labels分类完，如果同一个分类下还有不同种类，就以少数服从多数原则，指定多数为这个分类下的类型
        return majorityCnt(classList)
    # 获取当前根据属性划分，具有最大信息收益的属性
    bestFeat = chooseBestFeatureToSplit(dataSet)
    # 根据列号，获取对应列名(属性名)
    bestFeatLabel = labels[bestFeat]
    # 将该列名加入已被用来划分过的列表中
    featLabels.append(bestFeatLabel)
    # 创建最优特征的标签生成树 {'you_jiji':{ 0:{}, 1:{} }}
    myTree = {bestFeatLabel:{}}
    # 删除labels里面的第bestFeat个元素
    # 为什么要删，因为这里用了递归创建树，及create(dataSet, labels, featLabels)这个函数，故对应subData集的labels要小一点
    del(labels[bestFeat])
    featValues = [example[bestFeat] for example in dataSet]
    uniqueVals = set(featValues)
    # 用来递归生成决策树 like this: {'SX1': {0: 'c', 1: {'SX2': {0: 'f', 1: 'g'}}}}
    # value 在该特征是否有上取值,uniqueVals = [0,1] 0可表示为无脚蹼，1表示有
    for value in uniqueVals:
        # myTree-> '脚蹼' -0-> {}       bestFeatLabel  -value1-> {}
        #                -1-> {}                       -value2-> {}
        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), labels, featLabels)
    return myTree


"""
instructions: 返回该myTree树下的叶子数目, 返回树的层数
    有递归调用求叶子节点数
    next(iter(x))   # 用来遍历输出字典x的key的
    ** 有递归函数时，要先想该函数有这个功能。然后再去一步一步实现 **
        大体流程：
            1.当前分类的类名 like 有无jiji
            2.将该分类的树赋给secondDict    # 因为这个是字典树。{'SX1': {0: 'c', 1: {'SX2': {0: 'f', 1: 'g'}}}}
                要将secondDict = {0: 'c', 1: {'SX2': {0: 'f', 1: 'g'}}}
            3.遍历该树的每个分支，求其分支叶子树 + 1
                若该分支为叶子节点 return 1

Parameters: 
    myTree - 决策树(就一个字典树 {'SX1': {0: 'c', 1: {'SX2': {0: 'f', 1: 'g'}}}} )
Returns:
    numLeafs - 叶节点树
    maxDepth - 树的深度
Modify:
    2020-12-21
"""
def getNumLeafs(myTree):
    numLeafs = 0
    # 遍历当前层，就是每一层的key
    firstStr = next(iter(myTree))
    # 此时得到了一个分类下的树了. 比如这个 {0: 'c', 1: {'SX2': {0: 'f', 1: 'g'}}}
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[key]).__name__ == 'dict':
            numLeafs += getNumLeafs(secondDict[key])
        else:
            numLeafs += 1
    return numLeafs


def getTreeDepth(myTree):
    maxDepth = 0
    firstStr = next(iter(myTree))
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[key]).__name__ == 'dict':
            thisDepth = 1 + getTreeDepth(secondDict[key])
        else:
            thisDepth = 1
        
        if thisDepth > maxDepth:
            maxDepth = thisDepth
    
    return maxDepth


"""
instructions: 画出节点
    有两种节点 
        - 叶节点    leaf node
        - 决策节点  decision node
Parameters: 
     nodeTxt - 节点名称
     centerPt - 节点位置(及显示的文本"有脚蹼", 该值为文本框中心点的位置)
     parentPt - 箭头的启始位置(也可以认为是父节点的位置)
     nodeType - 节点类型
Returns:
    
Modify:
    2020-12-21 23

此部分与plotNode()单独连用
def createPlot():
    fig = plt.figure(1, facecolor="white")
    fig.clf()
    axprops = dict(xticks=[], yticks=[])
    createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)
    plotNode('decision_node', (0.5, 0.1), (0.1, 0.5), decisionNode)
    plotNode('leafs_node', (0.8, 0.1), (0.3, 0.8), leafNode)
    plt.show()
"""
# 生成字典 decisionNode = {'boxstyle':'sawtooth', 'fc':'0.8'}
# 此处定义节点的类型。波浪线框，直线框等
decisionNode = dict(boxstyle="sawtooth", fc='0.8')
leafNode = dict(boxstyle="round4", fc="0.8")

def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    arrow_args = dict(arrowstyle="<-")
    font = FontProperties(fname=r"./msyh.ttf", size=14)
    createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction', xytext=centerPt,
                            textcoords='axes fraction', va='center', ha='center', bbox=nodeType, 
                            arrowprops=arrow_args, FontProperties=font)


def plotMidText(cntrPt, parentPt, txtString):
    xMid = (parentPt[0]-cntrPt[0])/2.0 + cntrPt[0]
    yMid = (parentPt[1]-cntrPt[1])/2.0 + cntrPt[1]
    createPlot.ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)


"""
instructions: 将树画出来的函数
    该函数中主要疑惑点：
        1.xoff 偏移量怎么算的
        答：以一颗树为列，知道了当前有多少numleafs后，其每个叶节点的宽度可以确定，即为 width / numleafs，但是若以 1/w为第一个
            叶节点的x坐标，最后一个坐标为x = w/w即为1 (1/w,2/w,3/w,...,w/w)，则画出来整个树后，就会使树偏向figure的右边，如
            果左移1/w，则第一个叶节点x=0，又会偏向左边，因此需要将整个树左移1/(2*w),及不能移动1/w只能移动1/w的一半。
            * 由此x的起始坐标为：  -1/(2*w)
            * 相邻叶子节点的计算： 上一个叶节点 + 1/w

        2.cntrPt 怎么算的
        答：对于一颗树的根节点的x坐标： plotTree.xoff + (1 + float(numLeafs)) / 2.0 / plotTree.totalW，将其化开理解
            plotTree.xoff + float(numLeafs) / 2.0 / plotTree.totalW + 1 / 2.0 / plotTree.totalW，其中后半部分的理解
            为：有中间部分为得到叶节点后，根据叶节点的个数求得根节点在中心位置，但由于xoff初始为 -1/(2*w)，及向左偏移了这么多，故
            需要在这里将将这个偏移加回来
    
Parameters: 
     myTree - 要画的树
     parentPt - 文本框中心点位置
     cntrPt - 在这个函数里是变化的其作用，刚计算出来是文本框中心位置，画过一次后，成为下一层节点的parentPt
Returns:
    None
Modify:
    2020-12-23
"""
def plotTree(myTree, parentPt, nodeTxt):
    numLeafs = getNumLeafs(myTree)
    firstStr = next(iter(myTree))
    # 获取叶节点数后，相邻节点之间宽度为 l/w => 有n个叶子节点的根节点的x坐标为：(numleafs * l/w) / 2
    # cetrPt 记录的是当前树的根节点的坐标
    cntrPt = (plotTree.xoff + (1 + float(numLeafs)) / 2.0 / plotTree.totalW, plotTree.yoff)
    print("the parentPt is: " + str(parentPt))
    print("the cntrPt is: " + str(cntrPt))
    plotMidText(cntrPt, parentPt, nodeTxt)
    # firstStr 该文本框的字符
    # 并且这里画的是决策节点
    plotNode(firstStr, cntrPt, parentPt, decisionNode)
    secondDict = myTree[firstStr]
    plotTree.yoff = plotTree.yoff - 1.0 / plotTree.totalD   # y 每次减去 1/深度
    print("plotTree xoff is: " + str(plotTree.xoff))
    print("plotTree yoff is: " + str(plotTree.yoff))
    for key in secondDict.keys():
        if type(secondDict[key]).__name__ == 'dict':
            plotTree(secondDict[key], cntrPt, str(key))
        else:
            plotTree.xoff = plotTree.xoff + 1.0 / plotTree.totalW
            plotNode(secondDict[key], (plotTree.xoff, plotTree.yoff), cntrPt, leafNode)
            plotMidText((plotTree.xoff, plotTree.yoff), cntrPt, str(key))
    plotTree.yoff = plotTree.yoff + 1.0 / plotTree.totalD
    print("plotTree xoff is: " + str(plotTree.xoff))
    print("plotTree yoff is: " + str(plotTree.yoff))

"""
instructions: 创建树图
    创建图figure
    得到树的宽度，深度
    得到x,y的偏移量
    将树画出来(调用画树的函数)
    ** 需要注意的是:创建的figure的x轴范围：0.0～1.0 y轴的范围:0.0～1.0
Parameters: 
     inTree - 要画的树
Returns:
    None
Modify:
    2020-12-23
"""
def createPlot(inTree):
    fig = plt.figure(1, facecolor="white")
    fig.clf()
    axprops = dict(xticks=[], yticks=[])
    createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)
    plotTree.totalW = float(getNumLeafs(inTree))
    plotTree.totalD = float(getTreeDepth(inTree))
    print("the width is: " + str(plotTree.totalW))
    print("the depth is: " + str(plotTree.totalD))
    # xoff的偏移: 由 x轴长度/叶节点数w 1/w 为每个叶节点所占宽度, 然后1/(2*w)及为一个节点所占宽度的一半
    plotTree.xoff = -0.5 / plotTree.totalW
    plotTree.yoff = 1.0
    print("the xoff is: " + str(plotTree.xoff))
    print("the yoff is: " + str(plotTree.yoff))
    plotTree(inTree, (0.5, 1.0), '')        # 画树给的3个参数
    plt.show()


"""
instructions: 分类函数
    1.进入树的第一个分类节点
    2.比较测试向量对应列与当前在树中正在比较的节点
    3.进入对应左边树，或右边树中。
    4.如果还没分类完成，则继续分类。否则返回该节点的分类结果
Parameters: 
     inputTree - 创建好的树
     featLabels - 原数据集对应属性行
     testVec - 测试向量 [0,0,1,1]
Returns:
    None
Modify:
    2021-1-3
"""
def classify(inputTree, featLabels, testVec):
    firstStr = next(iter(inputTree))
    secondDict = inputTree[firstStr]
    # 通过特征名称获取特征索引值
    featIndex = featLabels.index(firstStr)
    # key 分别取0 ，1值进行分类。及进入左子树，还是右子树
    for key in secondDict.keys():
        if testVec[featIndex] == key:
            # 判断该给类下是否还有继续分类。及看其value是否还是一个字典，如果不是则直接返回value值（其实就是叶子节点）
            if type(secondDict[key]).__name__ == 'dict':
                classLabel = classify(secondDict[key], featLabels, testVec)
            else:
                classLabel = secondDict[key]
    return classLabel


"""
"""
def storeTree(inputTree, filename):
    with open(filename, 'wb') as fw:
        pickle.dump(inputTree, fw)


"""
"""
def grabTree(filename):
    fr = open(filename, 'rb')
    return pickle.load(fr)


def main():
    dataSet, labels = createDataSet()
    featLabels = []
    #print(calcShannonEnt(dataSet))
    #print(labels)
    #print(splitDataSet(dataSet, 1, 1))  # value = 0时,结果 [[1, 'B']]。 value = 1时，结果[[1, 'A'], [0, 'C'], [0, 'A']]。
    #print(chooseBestFeatureToSplit(dataSet))
    myTree = createTree(dataSet, labels, featLabels)
    storeTree(myTree, 'classifierStorage.txt')
    myTree1 = grabTree('classifierStorage.txt')
    print(myTree1)
    #print(tree)
    createPlot(myTree)

# 当在终端调用python3 trees.py时，就会自动调用该函数
if __name__ == '__main__':
    main()